open Base
open Unique_id.Ids

type 'a t =
  | Tree of Tree.tag * 'a t list * (Event.t -> 'a)
  | Dynamic of 'a t Behavior.t
  | Fold of Src.value * ('a -> Src.value -> Src.value) * (Src_id.t -> 'a t)
  | Bind of Var_id.t * Src_id.t * 'a t
  | Map: ('a -> 'b) * 'a t -> 'b t
and 'a component = 'a t

(*** constructors ***)

let rect w h =
  Tree(Tree.Rect(w, h), [], fun e -> `Evt e)

let arm dir cs =
  Tree(Tree.Arm(dir), cs, fun e -> `Evt e)

let fold ~init ~f make_comp =
  Fold(init, f, fun src -> make_comp (Behavior.of_src src))

let fold_var ~init ~f var comp =
  let init = Behavior.upcast var init in
  let f e x = Behavior.upcast1 var (f e) x in
  Fold(init, f, fun src -> Bind(Behavior.var_id var, src, comp))

let dyn b f =
  Dynamic(Behavior.map ~f b)

let rec map ~f = function
  | Map(g, c)  -> Map(Fn.compose f g, c)
  | Dynamic(b) -> Dynamic(Behavior.map ~f:(map ~f) b)
  | c          -> Map(f, c)

(*** environment ***)

type 'a env =
  { handlers : ('a -> Src.effect) list
  ; vars : Behavior.vars }

let map_env env ~f =
  { vars = env.vars
  ; handlers = List.map env.handlers
                 ~f:(fun h x -> h (f x)) }

let empty_env =
  { handlers = []
  ; vars = Var_id.empty_map }

let with_var env var src_id =
  { env with vars = Map.set env.vars ~key:var ~data:src_id }

let with_handler env h =
  { env with handlers = h :: env.handlers }

let bind_event_handlers env of_evt tree =
  Tree.set_event_handler tree @@ fun evt ->
  List.fold env.handlers
    ~init:Src.empty_effect
    ~f:(fun effect h ->
        Src.merge_effect effect (h (of_evt evt)))

(*** instantiation ***)

module type CONTEXT = sig
  type t
  type deleter

  val src_vals : t -> Src.values

  val create_source
    :  t
    -> init:Src.value
    -> Src_id.t

  val create_sink
    :  t
    -> tree:Tree.t
    -> env:'a env
    -> render:(Src.values -> 'a component)
    -> del:deleter
    -> Sink_id.t

  val connect
    :  t
    -> Sink_id.t
    -> Src_id.t
    -> unit

  val capture_deleter
    :  t
    -> f:(unit -> 'a)
    -> (deleter * 'a)
end

module type INST = sig
  type ctx
  val inst : ctx -> 'a env -> 'a component -> Tree.t
end

module Make_inst(Ctx : CONTEXT)
  : INST with type ctx := Ctx.t =
struct
  let rec inst : 'a. Ctx.t -> 'a env -> 'a component -> _ = fun ctx env -> function
    | Tree(tag, comps, of_evt) ->
       let tree = Tree.create tag in
       let () = Tree.set_children tree (List.map comps ~f:(inst ctx env)) in
       let () = bind_event_handlers env of_evt tree in
       tree

    | Dynamic(behav) ->
       let (deps, render) = behav.eval env.vars in
       let init_comp = render (Ctx.src_vals ctx) in
       let (del, tree) = Ctx.capture_deleter ctx
                           ~f:(fun () -> inst ctx env init_comp) in
       let sink_id = Ctx.create_sink ctx ~tree ~env ~render ~del in
       let () = Set.iter deps ~f:(Ctx.connect ctx sink_id) in
       tree

    | Fold(init, f, make_comp) ->
       let src = Ctx.create_source ctx ~init in
       inst ctx (with_handler env (fun evt -> Src.single_effect src (f evt)))
         (make_comp src)

    | Bind(var_id, src_id, comp) ->
       inst ctx (with_var env var_id src_id) comp

    | Map(f, comp) ->
       inst ctx (map_env ~f env) comp
end
